import math

import numpy as np
from sklearn import preprocessing
from metrics import *
import saliency_metrics

s2cnnPath = 'D:/VR_project/saliency-convnet-7.31/output_sal_maps/test/s2cnn_relu.npy'
attentionPath = 'D:/VR_project/saliency-convnet-7.31/sample_data/dataset2.npy'

s2cnn = np.load(s2cnnPath, allow_pickle=True)
atten = np.load(attentionPath, allow_pickle=True)
index = len(atten) - len(s2cnn)
fixation_maps = np.array([create_fixation_map(item[1]) for item in atten[index:]])
# fixation_maps = np.array([create_fixation_map(item[1]) for item in atten])
print(s2cnn.shape, fixation_maps.shape)

mmscaler = preprocessing.MinMaxScaler(feature_range=(0, 1))
saliency_maps = mmscaler.fit_transform(s2cnn.ravel().reshape(-1, 1)).reshape(s2cnn.shape)

CC_list = []
NSS_list = []
KL_list = []
AUC_list = []
sAUC_list = []

t = len(s2cnn)
# print(attentionList[i], fixation_maps.shape, saliency_maps.shape)
for idx in range(0, t):
    fixmap = fixation_maps[idx]
    salmap = saliency_maps[idx]
    if len(np.unique(fixmap)) != 2 or fixmap.shape != salmap.shape:
        continue
    CC = cal_CC(fixmap, salmap)
    NSS = cal_NSS(fixmap, salmap)
    KL = cal_KL(fixmap, salmap)
    AUC = cal_AUC(fixmap, salmap)
    sAUC = saliency_metrics.auc_shuff(salmap, fixmap, fixmap)

    if CC == -1 or NSS == -1 or math.isnan(CC):
        continue
    if KL == float('inf') or math.isnan(KL):
        continue
    CC_list.append(CC)
    NSS_list.append(NSS)
    KL_list.append(KL)
    AUC_list.append(AUC)
    sAUC_list.append(sAUC)
    print(f'\r[{idx}]/[{t}]', end=' ')

print(f'CC={round(np.mean(CC_list), 4)}, '
      f'NSS={round(np.mean(NSS_list), 4)}, '
      f'KL={round(np.mean(KL_list), 4)}, '
      f'AUC={round(np.mean(AUC_list), 4)} '
      f'sAUC={round(np.mean(sAUC_list), 4)}')